{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from typing import *\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from functools import partial\n",
    "from overrides import overrides\n",
    "\n",
    "from allennlp.data import Instance\n",
    "from allennlp.data.token_indexers import TokenIndexer\n",
    "from allennlp.data.tokenizers import Token\n",
    "from allennlp.nn import util as nn_util"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config(dict):\n",
    "    def __init__(self, **kwargs):\n",
    "        super().__init__(**kwargs)\n",
    "        for k, v in kwargs.items():\n",
    "            setattr(self, k, v)\n",
    "    \n",
    "    def set(self, key, val):\n",
    "        self[key] = val\n",
    "        setattr(self, key, val)\n",
    "        \n",
    "config = Config(\n",
    "    testing=True,\n",
    "    seed=1,\n",
    "    batch_size=64,\n",
    "    lr=3e-4,\n",
    "    epochs=2,\n",
    "    hidden_sz=64,\n",
    "    max_seq_len=100, # necessary to limit memory usage\n",
    "    max_vocab_size=100000,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.common.checks import ConfigurationError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_GPU = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = Path(\"../data\") / \"jigsaw\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set random seed manually to replicate results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x1061ae730>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.manual_seed(config.seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.vocabulary import Vocabulary\n",
    "from allennlp.data.dataset_readers import DatasetReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "label_cols = [\"toxic\", \"severe_toxic\", \"obscene\",\n",
    "              \"threat\", \"insult\", \"identity_hate\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.fields import TextField, MetadataField, ArrayField\n",
    "\n",
    "class JigsawDatasetReader(DatasetReader):\n",
    "    def __init__(self, tokenizer: Callable[[str], List[str]]=lambda x: x.split(),\n",
    "                 token_indexers: Dict[str, TokenIndexer] = None,\n",
    "                 max_seq_len: Optional[int]=config.max_seq_len) -> None:\n",
    "        super().__init__(lazy=False)\n",
    "        self.tokenizer = tokenizer\n",
    "        self.token_indexers = token_indexers or {\"tokens\": SingleIdTokenIndexer()}\n",
    "        self.max_seq_len = max_seq_len\n",
    "\n",
    "    @overrides\n",
    "    def text_to_instance(self, tokens: List[Token], id: str=None,\n",
    "                         labels: np.ndarray=None) -> Instance:\n",
    "        sentence_field = TextField(tokens, self.token_indexers)\n",
    "        fields = {\"tokens\": sentence_field}\n",
    "        \n",
    "        id_field = MetadataField(id)\n",
    "        fields[\"id\"] = id_field\n",
    "        \n",
    "        if labels is None:\n",
    "            labels = np.zeros(len(label_cols))\n",
    "        label_field = ArrayField(array=labels)\n",
    "        fields[\"label\"] = label_field\n",
    "\n",
    "        return Instance(fields)\n",
    "    \n",
    "    @overrides\n",
    "    def _read(self, file_path: str) -> Iterator[Instance]:\n",
    "        df = pd.read_csv(file_path)\n",
    "        if config.testing: df = df.head(1000)\n",
    "        for i, row in df.iterrows():\n",
    "            yield self.text_to_instance(\n",
    "                [Token(x) for x in self.tokenizer(row[\"comment_text\"])],\n",
    "                row[\"id\"], row[label_cols].values,\n",
    "            )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare token handlers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the spacy tokenizer here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.tokenizers.word_splitter import SpacyWordSplitter\n",
    "from allennlp.data.token_indexers.elmo_indexer import ELMoCharacterMapper, ELMoTokenCharactersIndexer\n",
    "\n",
    "# the token indexer is responsible for mapping tokens to integers\n",
    "token_indexer = ELMoTokenCharactersIndexer()\n",
    "\n",
    "def tokenizer(x: str):\n",
    "    return [w.text for w in\n",
    "            SpacyWordSplitter(language='en_core_web_sm', \n",
    "                              pos_tags=False).split_words(x)[:config.max_seq_len]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = JigsawDatasetReader(\n",
    "    tokenizer=tokenizer,\n",
    "    token_indexers={\"tokens\": token_indexer}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "267it [00:01, 78.76it/s]\n",
      "251it [00:00, 260.17it/s]\n"
     ]
    }
   ],
   "source": [
    "train_ds, test_ds = (reader.read(DATA_ROOT / fname) for fname in [\"train.csv\", \"test_proced.csv\"])\n",
    "val_ds = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "267"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<allennlp.data.instance.Instance at 0x1a22865208>,\n",
       " <allennlp.data.instance.Instance at 0x1a22848dd8>,\n",
       " <allennlp.data.instance.Instance at 0x1a2281fa58>,\n",
       " <allennlp.data.instance.Instance at 0x1a227c4c18>,\n",
       " <allennlp.data.instance.Instance at 0x1a227bc048>,\n",
       " <allennlp.data.instance.Instance at 0x1a227b16d8>,\n",
       " <allennlp.data.instance.Instance at 0x1a227ab198>,\n",
       " <allennlp.data.instance.Instance at 0x1a227c4588>,\n",
       " <allennlp.data.instance.Instance at 0x1a2058e5c0>,\n",
       " <allennlp.data.instance.Instance at 0x1a205af550>]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_ds[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tokens': [Explanation,\n",
       "  Why,\n",
       "  the,\n",
       "  edits,\n",
       "  made,\n",
       "  under,\n",
       "  my,\n",
       "  username,\n",
       "  Hardcore,\n",
       "  Metallica,\n",
       "  Fan,\n",
       "  were,\n",
       "  reverted,\n",
       "  ?,\n",
       "  They,\n",
       "  were,\n",
       "  n't,\n",
       "  vandalisms,\n",
       "  ,,\n",
       "  just,\n",
       "  closure,\n",
       "  on,\n",
       "  some,\n",
       "  GAs,\n",
       "  after,\n",
       "  I,\n",
       "  voted,\n",
       "  at,\n",
       "  New,\n",
       "  York,\n",
       "  Dolls,\n",
       "  FAC,\n",
       "  .,\n",
       "  And,\n",
       "  please,\n",
       "  do,\n",
       "  n't,\n",
       "  remove,\n",
       "  the,\n",
       "  template,\n",
       "  from,\n",
       "  the,\n",
       "  talk,\n",
       "  page,\n",
       "  since,\n",
       "  I,\n",
       "  'm,\n",
       "  retired,\n",
       "  now.89.205.38.27],\n",
       " '_token_indexers': {'tokens': <allennlp.data.token_indexers.elmo_indexer.ELMoTokenCharactersIndexer at 0x1a1f331860>},\n",
       " '_indexed_tokens': None,\n",
       " '_indexer_name_to_indexed_token': None}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vars(train_ds[0].fields[\"tokens\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare vocabulary"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We don't need to build the vocab: all that is handled by the token indexer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = Vocabulary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare iterator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The iterator is responsible for batching the data and preparing it for input into the model. We'll use the BucketIterator that batches text sequences of smilar lengths together."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.iterators import BucketIterator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterator = BucketIterator(batch_size=config.batch_size, \n",
    "                          sorting_keys=[(\"tokens\", \"num_tokens\")],\n",
    "                         )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to tell the iterator how to numericalize the text data. We do this by passing the vocabulary to the iterator. This step is easy to forget so be careful! "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "iterator.index_with(vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(iterator(train_ds)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tokens': {'tokens': tensor([[[259,  66, 111,  ..., 261, 261, 261],\n",
       "           [259,  45, 260,  ..., 261, 261, 261],\n",
       "           [259, 112, 103,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259, 117, 105,  ..., 261, 261, 261],\n",
       "           [259, 100, 112,  ..., 261, 261, 261],\n",
       "           [259, 101, 106,  ..., 261, 261, 261]],\n",
       "  \n",
       "          [[259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259,  85, 105,  ..., 261, 261, 261],\n",
       "           [259, 112, 115,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259,  47, 260,  ..., 261, 261, 261],\n",
       "           [259,  68, 109,  ..., 261, 261, 261],\n",
       "           [259,  99, 102,  ..., 261, 261, 261]],\n",
       "  \n",
       "          [[259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259,  35, 260,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259,  99, 102,  ..., 261, 261, 261],\n",
       "           [259, 100,  98,  ..., 261, 261, 261],\n",
       "           [259,  98, 260,  ..., 261, 261, 261]],\n",
       "  \n",
       "          ...,\n",
       "  \n",
       "          [[259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259,  62, 260,  ..., 261, 261, 261],\n",
       "           [259,  62, 260,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259,  98, 115,  ..., 261, 261, 261],\n",
       "           [259, 112, 103,  ..., 261, 261, 261],\n",
       "           [259, 101, 102,  ..., 261, 261, 261]],\n",
       "  \n",
       "          [[259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259,  67,  74,  ..., 261, 261, 261],\n",
       "           [259,  45, 260,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259,  35, 260,  ..., 261, 261, 261],\n",
       "           [259, 111, 112,  ..., 261, 261, 261],\n",
       "           [259, 102, 111,  ..., 261, 261, 261]],\n",
       "  \n",
       "          [[259,  69, 102,  ..., 261, 261, 261],\n",
       "           [259, 104, 115,  ..., 261, 261, 261],\n",
       "           [259, 110, 102,  ..., 261, 261, 261],\n",
       "           ...,\n",
       "           [259, 111, 112,  ..., 261, 261, 261],\n",
       "           [259, 116, 109,  ..., 261, 261, 261],\n",
       "           [259, 117, 105,  ..., 261, 261, 261]]])},\n",
       " 'id': ['006774d59329b7bd',\n",
       "  '0091aec11b57d12e',\n",
       "  '002746baedcdff10',\n",
       "  '00a6754c92b4af4f',\n",
       "  '0016e01b742b8da3',\n",
       "  '00584d887401f47b',\n",
       "  '006d11791d76b9f3',\n",
       "  '001cadfd324f8087',\n",
       "  '0063a8786a7034fc',\n",
       "  '001d8e7be417776a',\n",
       "  '009565ee1bc64e68'],\n",
       " 'label': tensor([[0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 1., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0.]])}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[259,  66, 111,  ..., 261, 261, 261],\n",
       "         [259,  45, 260,  ..., 261, 261, 261],\n",
       "         [259, 112, 103,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259, 117, 105,  ..., 261, 261, 261],\n",
       "         [259, 100, 112,  ..., 261, 261, 261],\n",
       "         [259, 101, 106,  ..., 261, 261, 261]],\n",
       "\n",
       "        [[259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259,  85, 105,  ..., 261, 261, 261],\n",
       "         [259, 112, 115,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259,  47, 260,  ..., 261, 261, 261],\n",
       "         [259,  68, 109,  ..., 261, 261, 261],\n",
       "         [259,  99, 102,  ..., 261, 261, 261]],\n",
       "\n",
       "        [[259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259,  35, 260,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259,  99, 102,  ..., 261, 261, 261],\n",
       "         [259, 100,  98,  ..., 261, 261, 261],\n",
       "         [259,  98, 260,  ..., 261, 261, 261]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259,  62, 260,  ..., 261, 261, 261],\n",
       "         [259,  62, 260,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259,  98, 115,  ..., 261, 261, 261],\n",
       "         [259, 112, 103,  ..., 261, 261, 261],\n",
       "         [259, 101, 102,  ..., 261, 261, 261]],\n",
       "\n",
       "        [[259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259,  67,  74,  ..., 261, 261, 261],\n",
       "         [259,  45, 260,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259,  35, 260,  ..., 261, 261, 261],\n",
       "         [259, 111, 112,  ..., 261, 261, 261],\n",
       "         [259, 102, 111,  ..., 261, 261, 261]],\n",
       "\n",
       "        [[259,  69, 102,  ..., 261, 261, 261],\n",
       "         [259, 104, 115,  ..., 261, 261, 261],\n",
       "         [259, 110, 102,  ..., 261, 261, 261],\n",
       "         ...,\n",
       "         [259, 111, 112,  ..., 261, 261, 261],\n",
       "         [259, 116, 109,  ..., 261, 261, 261],\n",
       "         [259, 117, 105,  ..., 261, 261, 261]]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[\"tokens\"][\"tokens\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([11, 100, 50])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[\"tokens\"][\"tokens\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.modules.seq2vec_encoders import Seq2VecEncoder, PytorchSeq2VecWrapper\n",
    "from allennlp.nn.util import get_text_field_mask\n",
    "from allennlp.models import Model\n",
    "from allennlp.modules.text_field_embedders import TextFieldEmbedder\n",
    "\n",
    "class BaselineModel(Model):\n",
    "    def __init__(self, word_embeddings: TextFieldEmbedder,\n",
    "                 encoder: Seq2VecEncoder,\n",
    "                 out_sz: int=len(label_cols)):\n",
    "        super().__init__(vocab)\n",
    "        self.word_embeddings = word_embeddings\n",
    "        self.encoder = encoder\n",
    "        self.projection = nn.Linear(self.encoder.get_output_dim(), out_sz)\n",
    "        self.loss = nn.BCEWithLogitsLoss()\n",
    "        \n",
    "    def forward(self, tokens: Dict[str, torch.Tensor],\n",
    "                id: Any, label: torch.Tensor) -> torch.Tensor:\n",
    "        mask = get_text_field_mask(tokens)\n",
    "        embeddings = self.word_embeddings(tokens)\n",
    "        state = self.encoder(embeddings, mask)\n",
    "        class_logits = self.projection(state)\n",
    "        \n",
    "        output = {\"class_logits\": class_logits}\n",
    "        output[\"loss\"] = self.loss(class_logits, label)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "02/07/2019 17:36:11 - INFO - allennlp.modules.elmo -   Initializing ELMo\n"
     ]
    }
   ],
   "source": [
    "from allennlp.modules.text_field_embedders import BasicTextFieldEmbedder\n",
    "from allennlp.modules.token_embedders import ElmoTokenEmbedder\n",
    "\n",
    "options_file = 'https://s3-us-west-2.amazonaws.com/allennlp/models/elmo/2x1024_128_2048cnn_1xhighway/elmo_2x1024_128_2048cnn_1xhighway_options.json'\n",
    "weight_file = 'https://s3-us-west-2.amazonaws.com/allennlp/models/elmo/2x1024_128_2048cnn_1xhighway/elmo_2x1024_128_2048cnn_1xhighway_weights.hdf5'\n",
    "\n",
    "elmo_embedder = ElmoTokenEmbedder(options_file, weight_file)\n",
    "word_embeddings = BasicTextFieldEmbedder({\"tokens\": elmo_embedder})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.modules.seq2vec_encoders import PytorchSeq2VecWrapper\n",
    "encoder: Seq2VecEncoder = PytorchSeq2VecWrapper(nn.LSTM(word_embeddings.get_output_dim(), config.hidden_sz, bidirectional=True, batch_first=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how simple and modular the code for initializing the model is. All the complexity is delegated to each component."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = BaselineModel(\n",
    "    word_embeddings, \n",
    "    encoder, \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_GPU: model.cuda()\n",
    "else: model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic sanity checks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = nn_util.move_to_device(batch, 0 if USE_GPU else -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokens = batch[\"tokens\"]\n",
    "labels = batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'tokens': tensor([[[259,  66, 111,  ..., 261, 261, 261],\n",
       "          [259,  45, 260,  ..., 261, 261, 261],\n",
       "          [259, 112, 103,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259, 117, 105,  ..., 261, 261, 261],\n",
       "          [259, 100, 112,  ..., 261, 261, 261],\n",
       "          [259, 101, 106,  ..., 261, 261, 261]],\n",
       " \n",
       "         [[259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259,  85, 105,  ..., 261, 261, 261],\n",
       "          [259, 112, 115,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259,  47, 260,  ..., 261, 261, 261],\n",
       "          [259,  68, 109,  ..., 261, 261, 261],\n",
       "          [259,  99, 102,  ..., 261, 261, 261]],\n",
       " \n",
       "         [[259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259,  35, 260,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259,  99, 102,  ..., 261, 261, 261],\n",
       "          [259, 100,  98,  ..., 261, 261, 261],\n",
       "          [259,  98, 260,  ..., 261, 261, 261]],\n",
       " \n",
       "         ...,\n",
       " \n",
       "         [[259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259,  62, 260,  ..., 261, 261, 261],\n",
       "          [259,  62, 260,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259,  98, 115,  ..., 261, 261, 261],\n",
       "          [259, 112, 103,  ..., 261, 261, 261],\n",
       "          [259, 101, 102,  ..., 261, 261, 261]],\n",
       " \n",
       "         [[259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259,  67,  74,  ..., 261, 261, 261],\n",
       "          [259,  45, 260,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259,  35, 260,  ..., 261, 261, 261],\n",
       "          [259, 111, 112,  ..., 261, 261, 261],\n",
       "          [259, 102, 111,  ..., 261, 261, 261]],\n",
       " \n",
       "         [[259,  69, 102,  ..., 261, 261, 261],\n",
       "          [259, 104, 115,  ..., 261, 261, 261],\n",
       "          [259, 110, 102,  ..., 261, 261, 261],\n",
       "          ...,\n",
       "          [259, 111, 112,  ..., 261, 261, 261],\n",
       "          [259, 116, 109,  ..., 261, 261, 261],\n",
       "          [259, 117, 105,  ..., 261, 261, 261]]])}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 1, 1,  ..., 1, 1, 1],\n",
       "        [1, 1, 1,  ..., 1, 1, 1],\n",
       "        [1, 1, 1,  ..., 1, 1, 1],\n",
       "        ...,\n",
       "        [1, 1, 1,  ..., 1, 1, 1],\n",
       "        [1, 1, 1,  ..., 1, 1, 1],\n",
       "        [1, 1, 1,  ..., 1, 1, 1]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mask = get_text_field_mask(tokens)\n",
    "mask"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0858,  0.0223, -0.1586,  0.1349, -0.0376, -0.0261],\n",
       "        [-0.1301,  0.0480,  0.0867,  0.0141,  0.2163, -0.1783],\n",
       "        [-0.0898,  0.0805,  0.1165, -0.1205, -0.0798, -0.0552],\n",
       "        [-0.1485,  0.0437,  0.3471, -0.0798,  0.0412, -0.0691],\n",
       "        [ 0.1237,  0.1616,  0.1916, -0.1183,  0.0581, -0.1586],\n",
       "        [ 0.2513,  0.0618,  0.1800,  0.1422,  0.3283, -0.1055],\n",
       "        [ 0.0311, -0.0209, -0.1196, -0.0920, -0.0282, -0.2505],\n",
       "        [-0.1946,  0.1308,  0.0727, -0.0361, -0.1095, -0.1665],\n",
       "        [-0.1263, -0.0665,  0.0816, -0.1481,  0.1190, -0.1457],\n",
       "        [-0.0762, -0.0154,  0.0472,  0.0230, -0.1451, -0.1299],\n",
       "        [-0.2562,  0.1345,  0.2318,  0.1293,  0.1449, -0.2039]],\n",
       "       grad_fn=<AddmmBackward>)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings = model.word_embeddings(tokens)\n",
    "state = model.encoder(embeddings, mask)\n",
    "class_logits = model.projection(state)\n",
    "class_logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'class_logits': tensor([[ 0.0389, -0.1613, -0.1306,  0.0497, -0.1056,  0.0422],\n",
       "         [-0.0693, -0.1457,  0.1586, -0.0689,  0.0846, -0.1528],\n",
       "         [-0.0561,  0.1929,  0.0745, -0.0527,  0.0945, -0.1080],\n",
       "         [-0.2267,  0.1510,  0.2946,  0.0169, -0.0651, -0.1007],\n",
       "         [ 0.0745,  0.0488,  0.2536, -0.1912,  0.0134, -0.0950],\n",
       "         [ 0.0587,  0.2195, -0.0123,  0.1830,  0.0590, -0.0560],\n",
       "         [ 0.0060, -0.1080, -0.1402, -0.1042,  0.0381, -0.1645],\n",
       "         [-0.1112,  0.1329,  0.1719,  0.1552, -0.0425, -0.1649],\n",
       "         [-0.1144,  0.2286,  0.1458, -0.1806,  0.1630, -0.2000],\n",
       "         [-0.0798, -0.0302, -0.0057,  0.0076,  0.0402, -0.0796],\n",
       "         [-0.1634,  0.1240,  0.1739,  0.1157,  0.0166, -0.0234]],\n",
       "        grad_fn=<AddmmBackward>),\n",
       " 'loss': tensor(0.6955, grad_fn=<BinaryCrossEntropyWithLogitsBackward>)}"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model(**batch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss = model(**batch)[\"loss\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(0.6934, grad_fn=<BinaryCrossEntropyWithLogitsBackward>)"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss.backward()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[-9.2858e-05, -4.9361e-05, -3.7288e-05,  ...,  3.5947e-06,\n",
       "           6.7686e-06,  5.7458e-05],\n",
       "         [-1.0111e-04,  1.9700e-04,  2.2716e-05,  ..., -1.1300e-04,\n",
       "          -2.2771e-05, -3.8708e-04],\n",
       "         [ 1.7676e-06,  1.4282e-06,  6.0405e-05,  ..., -3.9052e-05,\n",
       "           1.6350e-04,  1.5471e-04],\n",
       "         ...,\n",
       "         [-4.0847e-05,  4.4254e-05, -7.2378e-05,  ...,  2.0154e-04,\n",
       "           8.3410e-05, -1.2389e-04],\n",
       "         [-5.1695e-06, -4.6811e-06, -5.3316e-05,  ...,  1.1514e-05,\n",
       "          -1.1821e-04, -1.1049e-06],\n",
       "         [ 9.2464e-05,  4.6357e-05, -6.5559e-05,  ..., -1.3204e-04,\n",
       "          -2.1428e-04, -1.9575e-04]]),\n",
       " tensor([[-2.1633e-05,  8.7957e-07, -3.4178e-05,  ..., -9.0936e-06,\n",
       "           4.9037e-05,  1.9673e-05],\n",
       "         [ 4.6168e-05,  4.7835e-05,  6.2958e-05,  ..., -5.4513e-05,\n",
       "          -2.6419e-05,  1.8046e-06],\n",
       "         [ 2.0971e-05,  6.7690e-05, -7.4084e-05,  ..., -1.8694e-06,\n",
       "          -4.2899e-06,  1.2227e-05],\n",
       "         ...,\n",
       "         [-4.7440e-05, -1.9033e-04,  3.7619e-05,  ...,  2.4566e-04,\n",
       "          -1.8149e-05,  4.2449e-05],\n",
       "         [ 6.3985e-06,  1.1565e-05,  2.4952e-05,  ..., -1.8709e-05,\n",
       "          -3.9199e-05, -1.5171e-05],\n",
       "         [-2.8640e-05,  4.0967e-05,  5.0600e-05,  ..., -3.9222e-05,\n",
       "           4.0500e-06, -7.8656e-05]]),\n",
       " tensor([-1.8289e-04, -2.6908e-04, -5.1939e-04, -7.2012e-05,  4.8000e-04,\n",
       "         -3.1593e-04, -1.2665e-04,  1.3794e-04, -2.6114e-04,  2.8007e-04,\n",
       "          2.1538e-05,  1.2137e-03, -5.3436e-04, -9.8040e-04,  1.1253e-03,\n",
       "          6.8724e-05, -1.8145e-04,  2.1271e-04,  4.1548e-04,  1.5018e-04,\n",
       "         -2.1635e-04, -1.8048e-05, -4.0464e-05, -2.7178e-04,  1.5397e-05,\n",
       "          2.5399e-05, -1.9531e-04, -2.5291e-04,  1.6188e-04,  4.1508e-04,\n",
       "         -2.7233e-04,  2.4203e-04, -2.7715e-04,  1.7744e-04, -3.8623e-06,\n",
       "          3.9980e-05, -5.0311e-04,  1.4990e-05,  3.2888e-04,  7.8127e-05,\n",
       "         -5.8224e-04, -3.1990e-04, -4.6293e-06, -3.4443e-04, -1.4607e-04,\n",
       "          8.7014e-05, -1.0022e-05, -7.5172e-05, -1.5612e-03,  2.7014e-05,\n",
       "          1.6339e-04, -1.3362e-06, -5.8879e-05,  1.0310e-03,  2.8607e-04,\n",
       "          5.3653e-04, -1.8801e-04,  3.1148e-04,  9.8425e-05, -4.5256e-05,\n",
       "         -8.6616e-05,  4.1113e-04,  5.4935e-05,  1.0169e-04,  3.9380e-05,\n",
       "         -6.4266e-04, -2.5855e-04,  1.3207e-05,  2.1473e-04, -2.1532e-04,\n",
       "         -1.6356e-04,  2.1920e-04, -4.1288e-04,  2.2363e-04,  1.4386e-05,\n",
       "          4.0407e-04, -2.2122e-04, -1.3259e-03,  1.1047e-03, -7.3801e-05,\n",
       "         -1.9034e-04,  2.8585e-04,  6.8557e-04,  3.1578e-04, -5.8607e-05,\n",
       "         -2.4889e-05,  3.1479e-05, -4.6390e-04, -1.2204e-05,  8.2191e-05,\n",
       "         -3.5209e-04, -1.1277e-03,  7.5371e-05,  9.5581e-06, -1.4981e-04,\n",
       "          2.4498e-04, -3.8889e-04,  1.1692e-04, -1.2015e-04,  1.1258e-04,\n",
       "         -6.8194e-04,  2.6460e-04,  1.9391e-04, -1.4658e-04, -2.8695e-04,\n",
       "         -2.7773e-04,  7.4139e-05, -2.8039e-04, -7.7563e-05,  5.6996e-06,\n",
       "          3.6432e-06, -5.6888e-05, -1.3161e-03,  3.0357e-05, -1.5887e-04,\n",
       "          6.4288e-05,  1.9450e-04,  4.6303e-04,  2.2863e-04,  3.5556e-04,\n",
       "         -2.6141e-04,  4.2701e-04,  8.4789e-05, -3.1176e-06, -1.0718e-04,\n",
       "          4.4328e-05, -9.7466e-05,  3.6617e-04,  3.1579e-03,  4.8224e-03,\n",
       "         -1.6356e-03,  3.6920e-04, -2.6254e-03,  2.1345e-03, -5.7336e-03,\n",
       "         -1.9483e-03,  5.9474e-03, -1.7852e-03, -6.1445e-03, -4.4191e-03,\n",
       "         -1.0049e-02, -4.0846e-03, -9.2853e-03,  4.9335e-03, -1.0926e-03,\n",
       "          1.3147e-03, -2.8987e-03,  4.0139e-03, -2.9943e-03, -3.0421e-03,\n",
       "          7.1858e-03,  3.7457e-03,  2.7480e-04, -3.0752e-03,  2.6983e-03,\n",
       "         -6.9565e-03, -1.9007e-03, -2.7884e-03,  1.9140e-03, -1.1020e-03,\n",
       "          2.5542e-03,  1.1207e-03,  1.4054e-03, -1.9124e-04, -2.8571e-03,\n",
       "          1.0549e-03, -2.3865e-03, -1.8581e-03,  4.5822e-03, -3.1310e-03,\n",
       "          1.5557e-05,  1.5228e-03,  2.3360e-03, -8.8089e-04,  5.1199e-04,\n",
       "         -1.1610e-04, -6.0483e-03,  6.8591e-05, -2.7749e-03, -1.9582e-03,\n",
       "         -2.9191e-03,  5.7811e-03,  1.5078e-03,  1.2317e-03,  1.1030e-03,\n",
       "          3.0628e-03,  2.2106e-03,  1.5010e-03,  2.0449e-03,  3.3073e-03,\n",
       "         -2.3389e-03, -3.1255e-03, -4.0988e-04, -2.4041e-04, -5.1027e-04,\n",
       "         -9.1633e-05,  6.7165e-04, -4.3092e-04,  5.3300e-05,  2.0612e-04,\n",
       "         -1.3560e-04,  2.6270e-04, -1.9312e-04,  7.5937e-04, -2.6526e-04,\n",
       "         -1.0669e-03,  1.2701e-03, -1.5756e-04, -1.0190e-04,  2.9275e-04,\n",
       "          9.0443e-04,  1.6498e-04, -2.2931e-04,  5.3158e-05, -2.1521e-04,\n",
       "         -3.8123e-04, -1.4780e-05, -1.4956e-05, -3.7940e-04, -4.7910e-04,\n",
       "          1.6168e-04,  1.9206e-04, -1.1554e-04,  3.0967e-04, -4.3983e-04,\n",
       "          9.4490e-05, -6.3891e-05,  6.2164e-05, -9.7865e-04, -6.5990e-06,\n",
       "          4.2056e-04,  5.3304e-05, -6.9280e-04, -3.5037e-04,  8.9095e-05,\n",
       "         -4.4718e-04, -1.5922e-04,  6.0451e-05, -2.3626e-05, -1.2763e-04,\n",
       "         -1.8213e-03,  3.9953e-05,  1.4904e-04,  5.8299e-06,  9.8328e-06,\n",
       "          1.0570e-03,  3.6838e-04,  8.7990e-04, -8.2065e-05,  4.5531e-04,\n",
       "         -4.0542e-05, -9.9973e-05, -7.3163e-05,  6.6055e-04,  8.7863e-06,\n",
       "          2.0746e-04]),\n",
       " tensor([-1.8289e-04, -2.6908e-04, -5.1939e-04, -7.2012e-05,  4.8000e-04,\n",
       "         -3.1593e-04, -1.2665e-04,  1.3794e-04, -2.6114e-04,  2.8007e-04,\n",
       "          2.1538e-05,  1.2137e-03, -5.3436e-04, -9.8040e-04,  1.1253e-03,\n",
       "          6.8724e-05, -1.8145e-04,  2.1271e-04,  4.1548e-04,  1.5018e-04,\n",
       "         -2.1635e-04, -1.8048e-05, -4.0464e-05, -2.7178e-04,  1.5397e-05,\n",
       "          2.5399e-05, -1.9531e-04, -2.5291e-04,  1.6188e-04,  4.1508e-04,\n",
       "         -2.7233e-04,  2.4203e-04, -2.7715e-04,  1.7744e-04, -3.8623e-06,\n",
       "          3.9980e-05, -5.0311e-04,  1.4990e-05,  3.2888e-04,  7.8127e-05,\n",
       "         -5.8224e-04, -3.1990e-04, -4.6293e-06, -3.4443e-04, -1.4607e-04,\n",
       "          8.7014e-05, -1.0022e-05, -7.5172e-05, -1.5612e-03,  2.7014e-05,\n",
       "          1.6339e-04, -1.3362e-06, -5.8879e-05,  1.0310e-03,  2.8607e-04,\n",
       "          5.3653e-04, -1.8801e-04,  3.1148e-04,  9.8425e-05, -4.5256e-05,\n",
       "         -8.6616e-05,  4.1113e-04,  5.4935e-05,  1.0169e-04,  3.9380e-05,\n",
       "         -6.4266e-04, -2.5855e-04,  1.3207e-05,  2.1473e-04, -2.1532e-04,\n",
       "         -1.6356e-04,  2.1920e-04, -4.1288e-04,  2.2363e-04,  1.4386e-05,\n",
       "          4.0407e-04, -2.2122e-04, -1.3259e-03,  1.1047e-03, -7.3801e-05,\n",
       "         -1.9034e-04,  2.8585e-04,  6.8557e-04,  3.1578e-04, -5.8607e-05,\n",
       "         -2.4889e-05,  3.1479e-05, -4.6390e-04, -1.2204e-05,  8.2191e-05,\n",
       "         -3.5209e-04, -1.1277e-03,  7.5371e-05,  9.5581e-06, -1.4981e-04,\n",
       "          2.4498e-04, -3.8889e-04,  1.1692e-04, -1.2015e-04,  1.1258e-04,\n",
       "         -6.8194e-04,  2.6460e-04,  1.9391e-04, -1.4658e-04, -2.8695e-04,\n",
       "         -2.7773e-04,  7.4139e-05, -2.8039e-04, -7.7563e-05,  5.6996e-06,\n",
       "          3.6432e-06, -5.6888e-05, -1.3161e-03,  3.0357e-05, -1.5887e-04,\n",
       "          6.4288e-05,  1.9450e-04,  4.6303e-04,  2.2863e-04,  3.5556e-04,\n",
       "         -2.6141e-04,  4.2701e-04,  8.4789e-05, -3.1176e-06, -1.0718e-04,\n",
       "          4.4328e-05, -9.7466e-05,  3.6617e-04,  3.1579e-03,  4.8224e-03,\n",
       "         -1.6356e-03,  3.6920e-04, -2.6254e-03,  2.1345e-03, -5.7336e-03,\n",
       "         -1.9483e-03,  5.9474e-03, -1.7852e-03, -6.1445e-03, -4.4191e-03,\n",
       "         -1.0049e-02, -4.0846e-03, -9.2853e-03,  4.9335e-03, -1.0926e-03,\n",
       "          1.3147e-03, -2.8987e-03,  4.0139e-03, -2.9943e-03, -3.0421e-03,\n",
       "          7.1858e-03,  3.7457e-03,  2.7480e-04, -3.0752e-03,  2.6983e-03,\n",
       "         -6.9565e-03, -1.9007e-03, -2.7884e-03,  1.9140e-03, -1.1020e-03,\n",
       "          2.5542e-03,  1.1207e-03,  1.4054e-03, -1.9124e-04, -2.8571e-03,\n",
       "          1.0549e-03, -2.3865e-03, -1.8581e-03,  4.5822e-03, -3.1310e-03,\n",
       "          1.5557e-05,  1.5228e-03,  2.3360e-03, -8.8089e-04,  5.1199e-04,\n",
       "         -1.1610e-04, -6.0483e-03,  6.8591e-05, -2.7749e-03, -1.9582e-03,\n",
       "         -2.9191e-03,  5.7811e-03,  1.5078e-03,  1.2317e-03,  1.1030e-03,\n",
       "          3.0628e-03,  2.2106e-03,  1.5010e-03,  2.0449e-03,  3.3073e-03,\n",
       "         -2.3389e-03, -3.1255e-03, -4.0988e-04, -2.4041e-04, -5.1027e-04,\n",
       "         -9.1633e-05,  6.7165e-04, -4.3092e-04,  5.3300e-05,  2.0612e-04,\n",
       "         -1.3560e-04,  2.6270e-04, -1.9312e-04,  7.5937e-04, -2.6526e-04,\n",
       "         -1.0669e-03,  1.2701e-03, -1.5756e-04, -1.0190e-04,  2.9275e-04,\n",
       "          9.0443e-04,  1.6498e-04, -2.2931e-04,  5.3158e-05, -2.1521e-04,\n",
       "         -3.8123e-04, -1.4780e-05, -1.4956e-05, -3.7940e-04, -4.7910e-04,\n",
       "          1.6168e-04,  1.9206e-04, -1.1554e-04,  3.0967e-04, -4.3983e-04,\n",
       "          9.4490e-05, -6.3891e-05,  6.2164e-05, -9.7865e-04, -6.5990e-06,\n",
       "          4.2056e-04,  5.3304e-05, -6.9280e-04, -3.5037e-04,  8.9095e-05,\n",
       "         -4.4718e-04, -1.5922e-04,  6.0451e-05, -2.3626e-05, -1.2763e-04,\n",
       "         -1.8213e-03,  3.9953e-05,  1.4904e-04,  5.8299e-06,  9.8328e-06,\n",
       "          1.0570e-03,  3.6838e-04,  8.7990e-04, -8.2065e-05,  4.5531e-04,\n",
       "         -4.0542e-05, -9.9973e-05, -7.3163e-05,  6.6055e-04,  8.7863e-06,\n",
       "          2.0746e-04]),\n",
       " tensor([[-1.3693e-05,  3.5492e-05,  3.8836e-05,  ...,  1.5531e-05,\n",
       "          -2.5531e-05,  9.8788e-05],\n",
       "         [-1.0106e-04, -6.7122e-05, -4.3416e-05,  ...,  2.5320e-05,\n",
       "          -2.2854e-05, -1.2133e-05],\n",
       "         [-1.9106e-04, -2.8379e-04, -2.2337e-05,  ..., -2.5157e-05,\n",
       "          -6.3701e-05, -8.8909e-05],\n",
       "         ...,\n",
       "         [-1.3045e-04, -1.2610e-04,  1.5718e-04,  ..., -1.8985e-04,\n",
       "           1.5459e-05,  9.1975e-05],\n",
       "         [-1.3092e-06,  3.2634e-05,  5.0653e-05,  ..., -4.7065e-05,\n",
       "           1.4359e-05,  1.6218e-04],\n",
       "         [-1.3850e-04,  2.6071e-04,  9.0735e-05,  ..., -4.4709e-05,\n",
       "          -7.2825e-05, -4.0785e-05]]),\n",
       " tensor([[ 3.9198e-05,  3.6093e-05,  1.0073e-05,  ...,  3.2309e-05,\n",
       "           4.0834e-05,  1.8984e-05],\n",
       "         [ 5.2208e-06,  1.0176e-05, -1.4218e-05,  ...,  2.4117e-05,\n",
       "          -3.6023e-06,  2.0741e-05],\n",
       "         [-2.5719e-05,  3.7614e-05, -1.4760e-05,  ...,  2.9246e-06,\n",
       "           3.6332e-06,  3.3101e-05],\n",
       "         ...,\n",
       "         [ 4.2339e-05,  8.6734e-05, -2.8854e-05,  ...,  9.8323e-05,\n",
       "           6.3971e-05,  9.3885e-05],\n",
       "         [ 3.8755e-06, -6.2681e-06,  6.9113e-06,  ..., -1.6179e-06,\n",
       "          -2.7118e-05,  1.0559e-05],\n",
       "         [-8.8274e-06, -2.4860e-05, -2.3402e-05,  ...,  4.0792e-05,\n",
       "           5.9632e-05, -3.1582e-06]]),\n",
       " tensor([-3.1882e-04, -1.2575e-04, -7.0851e-05,  9.7861e-05,  8.0411e-04,\n",
       "         -6.0739e-05,  8.0672e-05, -9.2691e-05, -2.3453e-05, -1.1065e-05,\n",
       "          1.2653e-04,  2.3667e-05, -2.5144e-04, -4.7993e-04,  2.3883e-04,\n",
       "         -1.0091e-04,  3.8555e-06, -1.1296e-04, -1.6009e-04,  3.7939e-04,\n",
       "          2.4170e-04,  1.7414e-04, -1.2887e-04, -5.2342e-04, -6.8780e-07,\n",
       "         -2.4035e-04, -1.1058e-05, -2.4535e-04, -5.5690e-05,  2.5274e-04,\n",
       "          8.4641e-05, -9.5996e-05, -1.3207e-04,  7.8541e-06,  3.5860e-04,\n",
       "          3.0041e-04,  1.3250e-04, -5.9634e-04,  8.5756e-04, -9.9092e-05,\n",
       "          2.9627e-04, -5.6307e-04, -6.2134e-04, -1.0405e-04, -1.8999e-04,\n",
       "          3.8252e-04,  8.7723e-04, -9.6770e-04, -3.1937e-04,  1.1396e-03,\n",
       "         -9.7511e-05, -2.3939e-05, -9.3181e-05,  2.0967e-04,  1.5640e-04,\n",
       "         -3.6970e-04,  3.2044e-05,  3.9688e-04,  1.7632e-04,  2.9812e-04,\n",
       "          6.9706e-04, -5.2266e-04, -1.5035e-05, -1.0308e-04, -2.2488e-04,\n",
       "         -1.8326e-04, -2.1289e-04,  2.8298e-04,  2.1913e-04, -5.1340e-05,\n",
       "         -1.5224e-05, -3.8982e-05, -6.1014e-05, -2.1992e-06,  1.3516e-05,\n",
       "         -3.5956e-04, -2.3332e-04, -1.9069e-04,  2.3362e-04,  3.7090e-04,\n",
       "          5.3618e-05, -2.9890e-05, -1.4304e-04,  2.4078e-04,  4.3497e-04,\n",
       "          3.5448e-05,  2.3792e-04,  5.5625e-05, -2.8427e-05, -9.0276e-05,\n",
       "         -1.0241e-04, -2.7537e-04, -1.9759e-04,  6.5203e-05,  2.1828e-05,\n",
       "         -8.6435e-05, -2.1141e-04, -2.2201e-04,  3.2622e-04,  1.7075e-05,\n",
       "          3.4396e-04, -3.3821e-04,  6.6244e-04, -1.1198e-05,  5.1033e-04,\n",
       "         -5.7329e-04, -2.2965e-04, -7.4956e-05, -1.6131e-04,  1.2568e-04,\n",
       "          2.6790e-04, -5.4360e-04, -5.8064e-04,  9.4800e-04, -3.1586e-06,\n",
       "          3.3778e-05,  2.1468e-05, -1.6846e-04,  1.4202e-04,  2.5844e-05,\n",
       "          1.5071e-04,  7.7495e-05,  3.2375e-04,  2.5563e-04,  5.6895e-04,\n",
       "         -1.4890e-04,  1.4254e-04, -1.6291e-04,  1.6524e-03,  1.1015e-03,\n",
       "         -3.3566e-03, -2.3480e-03,  2.4269e-03, -7.8355e-06,  1.9382e-04,\n",
       "         -3.4848e-04,  8.8729e-04, -2.4050e-04,  1.7631e-03, -3.6839e-03,\n",
       "         -1.9833e-03,  1.9564e-03,  3.0658e-03, -4.9071e-03,  8.0148e-04,\n",
       "          1.1957e-03,  9.9154e-04, -3.9123e-03,  1.6994e-03, -3.6358e-03,\n",
       "         -1.2680e-03, -5.3975e-03, -8.2622e-04,  1.9391e-03,  1.8153e-03,\n",
       "          4.2626e-04,  1.5605e-03,  9.7955e-04,  2.4097e-03, -1.0837e-03,\n",
       "         -7.0819e-04, -2.2160e-03,  1.5242e-03,  1.8572e-03, -2.5351e-03,\n",
       "         -2.6722e-03,  4.7083e-03,  8.2712e-04,  4.1232e-03, -4.5716e-03,\n",
       "          2.0351e-03,  1.0964e-03,  1.0187e-03, -7.0810e-03,  3.2993e-03,\n",
       "         -4.5731e-03,  1.9813e-03, -4.1529e-03,  2.3412e-04, -3.7788e-04,\n",
       "         -7.6671e-04, -3.9124e-03, -7.5092e-04,  4.2677e-03, -1.8672e-03,\n",
       "         -1.1358e-02,  2.1387e-03, -1.5868e-03, -3.2008e-03,  2.4815e-03,\n",
       "          6.5891e-05,  7.7214e-04, -4.0723e-04, -7.8499e-05, -1.6537e-04,\n",
       "          1.1667e-04,  4.5361e-04, -2.3344e-04,  6.9808e-05, -1.2489e-04,\n",
       "         -6.0328e-05,  7.1351e-05,  9.5994e-05, -1.9108e-05, -3.3025e-04,\n",
       "         -4.8541e-04,  2.7484e-04, -1.8271e-04, -5.8893e-06, -1.0252e-04,\n",
       "         -2.2626e-04,  3.9077e-04,  2.7216e-04,  1.4024e-04, -1.8900e-04,\n",
       "         -5.5530e-04, -4.2707e-05, -2.9925e-04, -8.2783e-05, -4.5825e-04,\n",
       "         -3.9084e-05,  3.0867e-04, -5.2911e-05, -7.8673e-05, -1.4825e-04,\n",
       "         -2.0144e-04,  6.2794e-04,  4.0874e-04,  2.4637e-04, -8.1289e-04,\n",
       "          1.6562e-03, -1.2787e-04,  3.2519e-04, -4.7420e-04, -4.7531e-04,\n",
       "         -5.1285e-05, -2.5240e-04,  4.9952e-04,  7.5387e-04, -1.2315e-03,\n",
       "         -7.9711e-04,  1.3144e-03, -1.1580e-04, -2.4145e-05, -9.1914e-05,\n",
       "          5.8057e-05,  2.1737e-04, -3.0636e-04, -3.0068e-06,  1.2625e-03,\n",
       "          2.2120e-04,  4.6732e-04,  9.3160e-04, -6.1672e-04,  2.1024e-06,\n",
       "         -1.5065e-04]),\n",
       " tensor([-3.1882e-04, -1.2575e-04, -7.0851e-05,  9.7861e-05,  8.0411e-04,\n",
       "         -6.0739e-05,  8.0672e-05, -9.2691e-05, -2.3453e-05, -1.1065e-05,\n",
       "          1.2653e-04,  2.3667e-05, -2.5144e-04, -4.7993e-04,  2.3883e-04,\n",
       "         -1.0091e-04,  3.8555e-06, -1.1296e-04, -1.6009e-04,  3.7939e-04,\n",
       "          2.4170e-04,  1.7414e-04, -1.2887e-04, -5.2342e-04, -6.8780e-07,\n",
       "         -2.4035e-04, -1.1058e-05, -2.4535e-04, -5.5690e-05,  2.5274e-04,\n",
       "          8.4641e-05, -9.5996e-05, -1.3207e-04,  7.8541e-06,  3.5860e-04,\n",
       "          3.0041e-04,  1.3250e-04, -5.9634e-04,  8.5756e-04, -9.9092e-05,\n",
       "          2.9627e-04, -5.6307e-04, -6.2134e-04, -1.0405e-04, -1.8999e-04,\n",
       "          3.8252e-04,  8.7723e-04, -9.6770e-04, -3.1937e-04,  1.1396e-03,\n",
       "         -9.7511e-05, -2.3939e-05, -9.3181e-05,  2.0967e-04,  1.5640e-04,\n",
       "         -3.6970e-04,  3.2044e-05,  3.9688e-04,  1.7632e-04,  2.9812e-04,\n",
       "          6.9706e-04, -5.2266e-04, -1.5035e-05, -1.0308e-04, -2.2488e-04,\n",
       "         -1.8326e-04, -2.1289e-04,  2.8298e-04,  2.1913e-04, -5.1340e-05,\n",
       "         -1.5224e-05, -3.8982e-05, -6.1014e-05, -2.1992e-06,  1.3516e-05,\n",
       "         -3.5956e-04, -2.3332e-04, -1.9069e-04,  2.3362e-04,  3.7090e-04,\n",
       "          5.3618e-05, -2.9890e-05, -1.4304e-04,  2.4078e-04,  4.3497e-04,\n",
       "          3.5448e-05,  2.3792e-04,  5.5625e-05, -2.8427e-05, -9.0276e-05,\n",
       "         -1.0241e-04, -2.7537e-04, -1.9759e-04,  6.5203e-05,  2.1828e-05,\n",
       "         -8.6435e-05, -2.1141e-04, -2.2201e-04,  3.2622e-04,  1.7075e-05,\n",
       "          3.4396e-04, -3.3821e-04,  6.6244e-04, -1.1198e-05,  5.1033e-04,\n",
       "         -5.7329e-04, -2.2965e-04, -7.4956e-05, -1.6131e-04,  1.2568e-04,\n",
       "          2.6790e-04, -5.4360e-04, -5.8064e-04,  9.4800e-04, -3.1586e-06,\n",
       "          3.3778e-05,  2.1468e-05, -1.6846e-04,  1.4202e-04,  2.5844e-05,\n",
       "          1.5071e-04,  7.7495e-05,  3.2375e-04,  2.5563e-04,  5.6895e-04,\n",
       "         -1.4890e-04,  1.4254e-04, -1.6291e-04,  1.6524e-03,  1.1015e-03,\n",
       "         -3.3566e-03, -2.3480e-03,  2.4269e-03, -7.8355e-06,  1.9382e-04,\n",
       "         -3.4848e-04,  8.8729e-04, -2.4050e-04,  1.7631e-03, -3.6839e-03,\n",
       "         -1.9833e-03,  1.9564e-03,  3.0658e-03, -4.9071e-03,  8.0148e-04,\n",
       "          1.1957e-03,  9.9154e-04, -3.9123e-03,  1.6994e-03, -3.6358e-03,\n",
       "         -1.2680e-03, -5.3975e-03, -8.2622e-04,  1.9391e-03,  1.8153e-03,\n",
       "          4.2626e-04,  1.5605e-03,  9.7955e-04,  2.4097e-03, -1.0837e-03,\n",
       "         -7.0819e-04, -2.2160e-03,  1.5242e-03,  1.8572e-03, -2.5351e-03,\n",
       "         -2.6722e-03,  4.7083e-03,  8.2712e-04,  4.1232e-03, -4.5716e-03,\n",
       "          2.0351e-03,  1.0964e-03,  1.0187e-03, -7.0810e-03,  3.2993e-03,\n",
       "         -4.5731e-03,  1.9813e-03, -4.1529e-03,  2.3412e-04, -3.7788e-04,\n",
       "         -7.6671e-04, -3.9124e-03, -7.5092e-04,  4.2677e-03, -1.8672e-03,\n",
       "         -1.1358e-02,  2.1387e-03, -1.5868e-03, -3.2008e-03,  2.4815e-03,\n",
       "          6.5891e-05,  7.7214e-04, -4.0723e-04, -7.8499e-05, -1.6537e-04,\n",
       "          1.1667e-04,  4.5361e-04, -2.3344e-04,  6.9808e-05, -1.2489e-04,\n",
       "         -6.0328e-05,  7.1351e-05,  9.5994e-05, -1.9108e-05, -3.3025e-04,\n",
       "         -4.8541e-04,  2.7484e-04, -1.8271e-04, -5.8893e-06, -1.0252e-04,\n",
       "         -2.2626e-04,  3.9077e-04,  2.7216e-04,  1.4024e-04, -1.8900e-04,\n",
       "         -5.5530e-04, -4.2707e-05, -2.9925e-04, -8.2783e-05, -4.5825e-04,\n",
       "         -3.9084e-05,  3.0867e-04, -5.2911e-05, -7.8673e-05, -1.4825e-04,\n",
       "         -2.0144e-04,  6.2794e-04,  4.0874e-04,  2.4637e-04, -8.1289e-04,\n",
       "          1.6562e-03, -1.2787e-04,  3.2519e-04, -4.7420e-04, -4.7531e-04,\n",
       "         -5.1285e-05, -2.5240e-04,  4.9952e-04,  7.5387e-04, -1.2315e-03,\n",
       "         -7.9711e-04,  1.3144e-03, -1.1580e-04, -2.4145e-05, -9.1914e-05,\n",
       "          5.8057e-05,  2.1737e-04, -3.0636e-04, -3.0068e-06,  1.2625e-03,\n",
       "          2.2120e-04,  4.6732e-04,  9.3160e-04, -6.1672e-04,  2.1024e-06,\n",
       "         -1.5065e-04])]"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[x.grad for x in list(model.encoder.parameters())]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=config.lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.training.trainer import Trainer\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    optimizer=optimizer,\n",
    "    iterator=iterator,\n",
    "    train_dataset=train_ds,\n",
    "    cuda_device=0 if USE_GPU else -1,\n",
    "    num_epochs=config.epochs,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "02/07/2019 17:36:24 - INFO - allennlp.training.trainer -   Beginning training.\n",
      "02/07/2019 17:36:24 - INFO - allennlp.training.trainer -   Epoch 0/1\n",
      "02/07/2019 17:36:24 - INFO - allennlp.training.trainer -   Peak CPU memory usage MB: 1255.145472\n",
      "02/07/2019 17:36:24 - INFO - allennlp.training.trainer -   Training\n",
      "loss: 0.6855 ||: 100%|██████████| 5/5 [00:36<00:00,  8.01s/it]\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -                     Training |  Validation\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   loss          |     0.686  |       N/A\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   cpu_memory_MB |  1255.145  |       N/A\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   Epoch duration: 00:00:36\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   Estimated training time remaining: 0:00:36\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   Epoch 1/1\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   Peak CPU memory usage MB: 3154.649088\n",
      "02/07/2019 17:37:01 - INFO - allennlp.training.trainer -   Training\n",
      "loss: 0.6464 ||: 100%|██████████| 5/5 [00:51<00:00, 11.81s/it]\n",
      "02/07/2019 17:37:52 - INFO - allennlp.training.trainer -                     Training |  Validation\n",
      "02/07/2019 17:37:52 - INFO - allennlp.training.trainer -   loss          |     0.646  |       N/A\n",
      "02/07/2019 17:37:52 - INFO - allennlp.training.trainer -   cpu_memory_MB |  3154.649  |       N/A\n",
      "02/07/2019 17:37:52 - INFO - allennlp.training.trainer -   Epoch duration: 00:00:51\n"
     ]
    }
   ],
   "source": [
    "metrics = trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generating Predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.iterators import DataIterator\n",
    "from tqdm import tqdm\n",
    "from scipy.special import expit # the sigmoid function\n",
    "\n",
    "def tonp(tsr): return tsr.detach().cpu().numpy()\n",
    "\n",
    "class Predictor:\n",
    "    def __init__(self, model: Model, iterator: DataIterator,\n",
    "                 cuda_device: int=-1) -> None:\n",
    "        self.model = model\n",
    "        self.iterator = iterator\n",
    "        self.cuda_device = cuda_device\n",
    "        \n",
    "    def _extract_data(self, batch) -> np.ndarray:\n",
    "        out_dict = self.model(**batch)\n",
    "        return expit(tonp(out_dict[\"class_logits\"]))\n",
    "    \n",
    "    def predict(self, ds: Iterable[Instance]) -> np.ndarray:\n",
    "        pred_generator = self.iterator(ds, num_epochs=1, shuffle=False)\n",
    "        self.model.eval()\n",
    "        pred_generator_tqdm = tqdm(pred_generator,\n",
    "                                   total=self.iterator.get_num_batches(ds))\n",
    "        preds = []\n",
    "        with torch.no_grad():\n",
    "            for batch in pred_generator_tqdm:\n",
    "                batch = nn_util.move_to_device(batch, self.cuda_device)\n",
    "                preds.append(self._extract_data(batch))\n",
    "        return np.concatenate(preds, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "from allennlp.data.iterators import BasicIterator\n",
    "# iterate over the dataset without changing its order\n",
    "seq_iterator = BasicIterator(batch_size=64)\n",
    "seq_iterator.index_with(vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [01:44<00:00, 18.25s/it]\n",
      "100%|██████████| 4/4 [00:50<00:00, 13.33s/it]\n"
     ]
    }
   ],
   "source": [
    "predictor = Predictor(model, seq_iterator, cuda_device=0 if USE_GPU else -1)\n",
    "train_preds = predictor.predict(train_ds) \n",
    "test_preds = predictor.predict(test_ds) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
